import argparse
import pandas as pd
import json
import os
import ast
import numpy as np

def extract_first_element(predictions_str):
    try:
        parsed = ast.literal_eval(predictions_str)
        return parsed[0][0][0]
    except (ValueError, IndexError, SyntaxError) as e:
        return predictions_str[0]

def parquet_to_json(file_path):
    df = pd.read_parquet(file_path, engine='auto')
    columns=['full_prompt', 'gold', 'metrics',
       'predictions', 'specifics', 'truncated']
    df = df[columns]
    df['predictions'] = df['predictions'].apply(extract_first_element)
    json_file_path = os.path.splitext(file_path)[0] + '.json'
    
    with open(json_file_path, 'w') as f:
        def default_serializer(obj):
            if isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, np.generic):
                return obj.item()
            raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
        json.dump(df.to_dict(orient='records'), f, indent=4, default=default_serializer)    
    
    return json_file_path

def json_to_jsonl(json_file_path):

    with open(json_file_path, 'r') as infile:
        data = json.load(infile)

    jsonl_file_path = os.path.splitext(json_file_path)[0] + '.jsonl'
    with open(jsonl_file_path, 'w') as outfile:
        for item in data:
            json.dump(item, outfile)
            outfile.write('\n')


def main(data_files):
    for file in data_files:
        json_file = parquet_to_json(file)
        json_to_jsonl(json_file)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Parquet to JSON,JSONL")
    parser.add_argument('--data_files', nargs='+', required=True, help='Parquet paths')
    args = parser.parse_args()

    main(args.data_files)